"""
批量缺陷检测脚本 - 处理整个数据集
支持断点续传、进度显示、统计报告
"""

import os
import json
import time
from pathlib import Path
from typing import List, Dict

import cv2
import google.generativeai as genai
import numpy as np
import PIL.Image

# 从 annotate.py 导入处理函数
from annotate import (
    CONFIG, 
    preprocess_image, 
    call_model, 
    filter_defects, 
    merge_nearby_defects,
    save_yolo,
    draw_boxes
)


# 批量处理配置
BATCH_CONFIG = {
    "input_dir": "风机叶片数据集/images",  # 输入图片目录
    "output_dir": "风机叶片数据集/labels",  # YOLO标注输出目录
    "visualization_dir": "风机叶片数据集/visualizations",  # 可视化图片输出目录
    "skip_existing": True,  # 跳过已标注的图片
    "save_visualization": True,  # 保存可视化图片
    "delay_between_requests": 2,  # API请求间隔（秒）
    "save_summary": True,  # 保存统计报告
}


class BatchProcessor:
    """批量处理器"""
    
    def __init__(self):
        self.input_dir = Path(BATCH_CONFIG["input_dir"])
        self.output_dir = Path(BATCH_CONFIG["output_dir"])
        self.viz_dir = Path(BATCH_CONFIG["visualization_dir"])
        
        # 创建输出目录
        self.output_dir.mkdir(parents=True, exist_ok=True)
        if BATCH_CONFIG["save_visualization"]:
            self.viz_dir.mkdir(parents=True, exist_ok=True)
        
        # 统计信息
        self.stats = {
            "total": 0,
            "processed": 0,
            "skipped": 0,
            "failed": 0,
            "no_defects": 0,
            "total_defects": 0,
            "defects_by_type": {cls: 0 for cls in CONFIG["defect_classes"]},
            "failed_images": []
        }
    
    def get_image_list(self) -> List[Path]:
        """获取待处理的图片列表"""
        extensions = ["*.jpg", "*.jpeg", "*.png", "*.JPG", "*.JPEG", "*.PNG"]
        image_files = []
        
        for ext in extensions:
            image_files.extend(self.input_dir.glob(ext))
        
        image_files = sorted(image_files)  # 排序以便确定性处理
        
        # 跳过已标注的图片
        if BATCH_CONFIG["skip_existing"]:
            filtered = []
            for img_path in image_files:
                label_path = self.output_dir / f"{img_path.stem}.txt"
                if not label_path.exists():
                    filtered.append(img_path)
                else:
                    self.stats["skipped"] += 1
            
            return filtered
        
        return image_files
    
    def process_single_image(self, image_path: Path) -> bool:
        """处理单张图片"""
        try:
            print(f"\n{'='*70}")
            print(f"📷 {image_path.name}")
            
            # 1. 读取图片
            image = PIL.Image.open(image_path)
            width, height = image.size
            print(f"   尺寸: {width}x{height}")
            
            # 2. 预处理
            print(f"   [1/4] 图像预处理...")
            processed = preprocess_image(image)
            
            # 3. AI检测
            print(f"   [2/4] AI模型分析...")
            raw_results = call_model(processed)
            
            if raw_results is None:
                print(f"   ❌ API调用失败")
                self.stats["failed"] += 1
                self.stats["failed_images"].append(str(image_path.name))
                return False
            
            # 4. 过滤和合并
            print(f"   [3/4] 结果过滤（原始: {len(raw_results)} 个）")
            defects = filter_defects(raw_results)
            
            if not defects:
                print(f"   ℹ️  未检测到有效缺陷")
                # 创建空标注文件
                label_path = self.output_dir / f"{image_path.stem}.txt"
                label_path.write_text("", encoding="utf-8")
                
                # 即使没有缺陷，也生成原图副本作为可视化
                if BATCH_CONFIG["save_visualization"]:
                    import shutil
                    viz_path = self.viz_dir / f"{image_path.stem}_annotated.jpg"
                    shutil.copy(str(image_path), str(viz_path))
                    print(f"   📄 已保存原图副本")
                
                self.stats["no_defects"] += 1
                self.stats["processed"] += 1
                return True
            
            if CONFIG["merge_nearby_defects"]:
                before_merge = len(defects)
                defects = merge_nearby_defects(defects)
                print(f"   [4/4] 缺陷合并: {before_merge} → {len(defects)} 个")
            else:
                print(f"   [4/4] 跳过合并: {len(defects)} 个")
            
            # 5. 统计
            self.stats["total_defects"] += len(defects)
            for defect in defects:
                dtype = defect["defect_type"]
                if dtype in self.stats["defects_by_type"]:
                    self.stats["defects_by_type"][dtype] += 1
            
            # 显示检测结果
            counts = {}
            for d in defects:
                counts[d["defect_type"]] = counts.get(d["defect_type"], 0) + 1
            print(f"   ✅ 检测完成:")
            for dtype, count in sorted(counts.items()):
                print(f"      - {dtype}: {count} 个")
            
            # 6. 保存YOLO标注
            label_path = self.output_dir / f"{image_path.stem}.txt"
            self._save_yolo_annotation(defects, label_path)
            
            # 7. 保存可视化
            if BATCH_CONFIG["save_visualization"]:
                viz_path = self.viz_dir / f"{image_path.stem}_annotated.jpg"
                self._save_visualization(defects, image_path, viz_path)
            
            self.stats["processed"] += 1
            return True
            
        except Exception as e:
            print(f"   ❌ 处理异常: {e}")
            self.stats["failed"] += 1
            self.stats["failed_images"].append(str(image_path.name))
            import traceback
            traceback.print_exc()
            return False
    
    def _save_yolo_annotation(self, defects: List[Dict], output_path: Path):
        """保存YOLO格式标注"""
        lines = []
        for d in defects:
            idx = CONFIG["defect_classes"].index(d["defect_type"])
            b = d["bounding_box"]
            xc, yc = (b["x1"] + b["x2"]) / 2, (b["y1"] + b["y2"]) / 2
            w, h = b["x2"] - b["x1"], b["y2"] - b["y1"]
            lines.append(f"{idx} {xc:.6f} {yc:.6f} {w:.6f} {h:.6f}")
        
        output_path.write_text("\n".join(lines), encoding="utf-8")
    
    def _save_visualization(self, defects: List[Dict], image_path: Path, output_path: Path):
        """保存可视化图片"""
        img = cv2.imread(str(image_path))
        if img is None:
            return
        
        h, w = img.shape[:2]
        colors = {
            "Corrosion": (0, 0, 255),
            "Crack": (255, 0, 0),
            "Peeled Paint": (0, 255, 255),
            "Damage": (0, 165, 255),
            "Coating Degradation": (0, 255, 0)
        }
        
        for d in defects:
            b = d["bounding_box"]
            x1, y1 = int(b["x1"] * w), int(b["y1"] * h)
            x2, y2 = int(b["x2"] * w), int(b["y2"] * h)
            color = colors.get(d["defect_type"], (255, 255, 255))
            
            # 绘制边框（更粗更明显）
            cv2.rectangle(img, (x1, y1), (x2, y2), color, 3)  # 3像素粗
            
            # 绘制标签背景和文字（更大更清晰）
            label = f"{d['defect_type']} {d['confidence']:.2f}"
            (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)  # 字体0.6，粗细2
            cv2.rectangle(img, (x1, y1 - th - 8), (x1 + tw + 6, y1), color, -1)
            cv2.putText(img, label, (x1 + 3, y1 - 4), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
        
        cv2.imwrite(str(output_path), img)
    
    def run(self):
        """执行批量处理"""
        print("\n" + "="*70)
        print("🚀 风力发电叶片缺陷检测 - 批量处理系统")
        print("="*70)
        print(f"📁 输入目录: {self.input_dir}")
        print(f"📁 标注输出: {self.output_dir}")
        if BATCH_CONFIG["save_visualization"]:
            print(f"📁 可视化输出: {self.viz_dir}")
        print(f"⏱️  请求间隔: {BATCH_CONFIG['delay_between_requests']} 秒")
        
        # 获取图片列表
        image_list = self.get_image_list()
        self.stats["total"] = len(image_list) + self.stats["skipped"]
        
        if self.stats["skipped"] > 0:
            print(f"⏭️  跳过已标注: {self.stats['skipped']} 张")
        
        if not image_list:
            print("\n⚠️  没有需要处理的新图片")
            return
        
        print(f"\n📊 待处理: {len(image_list)} 张图片")
        print("="*70)
        
        # 开始处理
        start_time = time.time()
        
        for idx, image_path in enumerate(image_list, 1):
            print(f"\n[{idx}/{len(image_list)}] 进度: {idx/len(image_list)*100:.1f}%")
            
            self.process_single_image(image_path)
            
            # API限流
            if idx < len(image_list):
                time.sleep(BATCH_CONFIG["delay_between_requests"])
        
        elapsed_time = time.time() - start_time
        
        # 显示摘要
        self._print_summary(elapsed_time)
        
        # 保存报告
        if BATCH_CONFIG["save_summary"]:
            self._save_summary_report(elapsed_time)
    
    def _print_summary(self, elapsed_time: float):
        """打印处理摘要"""
        print("\n" + "="*70)
        print("📊 批量处理完成 - 统计报告")
        print("="*70)
        print(f"\n图片处理:")
        print(f"  总计: {self.stats['total']} 张")
        print(f"  成功处理: {self.stats['processed']} 张")
        print(f"  跳过已标注: {self.stats['skipped']} 张")
        print(f"  处理失败: {self.stats['failed']} 张")
        print(f"  无缺陷: {self.stats['no_defects']} 张")
        
        print(f"\n缺陷统计:")
        print(f"  总缺陷数: {self.stats['total_defects']} 个")
        if self.stats['total_defects'] > 0:
            print(f"  类型分布:")
            for dtype, count in sorted(self.stats["defects_by_type"].items()):
                if count > 0:
                    percentage = count / self.stats['total_defects'] * 100
                    bar = "█" * int(percentage / 5)
                    print(f"    {dtype:20s}: {count:3d} 个 ({percentage:5.1f}%) {bar}")
        
        print(f"\n性能指标:")
        print(f"  总耗时: {elapsed_time:.1f} 秒 ({elapsed_time/60:.1f} 分钟)")
        if self.stats['processed'] > 0:
            print(f"  平均速度: {elapsed_time/self.stats['processed']:.1f} 秒/张")
        
        if self.stats['failed'] > 0:
            print(f"\n⚠️  失败的图片:")
            for img_name in self.stats['failed_images']:
                print(f"    - {img_name}")
        
        print("="*70 + "\n")
    
    def _save_summary_report(self, elapsed_time: float):
        """保存JSON格式的统计报告"""
        report = {
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            "processing_stats": {
                "total_images": self.stats["total"],
                "processed": self.stats["processed"],
                "skipped": self.stats["skipped"],
                "failed": self.stats["failed"],
                "no_defects": self.stats["no_defects"],
            },
            "defect_stats": {
                "total_defects": self.stats["total_defects"],
                "by_type": self.stats["defects_by_type"],
            },
            "performance": {
                "elapsed_time_seconds": elapsed_time,
                "average_time_per_image": elapsed_time / max(self.stats['processed'], 1),
            },
            "failed_images": self.stats["failed_images"],
            "config": {
                "model": CONFIG["model_name"],
                "min_confidence": CONFIG["min_confidence"],
                "min_box_area": CONFIG["min_box_area"],
                "merge_nearby_defects": CONFIG["merge_nearby_defects"],
                "merge_distance_threshold": CONFIG["merge_distance_threshold"],
            }
        }
        
        report_path = self.output_dir / f"batch_report_{time.strftime('%Y%m%d_%H%M%S')}.json"
        report_path.write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8")
        
        print(f"📄 详细报告已保存: {report_path}")


def main():
    """主函数"""
    # 检查环境变量
    if not os.environ.get("GOOGLE_API_KEY"):
        print("❌ 未检测到环境变量 GOOGLE_API_KEY")
        print("请设置: set GOOGLE_API_KEY=your_api_key")
        return
    
    # 配置API
    genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
    
    try:
        processor = BatchProcessor()
        processor.run()
    except KeyboardInterrupt:
        print("\n\n⚠️  用户中断，正在安全退出...")
    except Exception as e:
        print(f"\n❌ 批量处理出错: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()

